from absl import app, flags

import collections
import pandas
import functools
import logging
import sys

import matplotlib
import matplotlib.pyplot as plt

from coin.base.book.types import TradeSide
from coin.base.datetime_util import to_timestamp_int

import coin.strategy.mm.feed as feed
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
from coin.strategy.mm.tool.util import get_order_sign, KuonaOrderManager

FLAGS = flags.FLAGS

import coin.strategy.mm.tool.plot_util as putil

import coin.strategy.mm.tool.archive_base as abase
from coin.strategy.accounting.pnl_balance_printer.util import is_inverse
from coin.support.pta.logic.pta_stats_calculator import get_futures_multiplier


def get_multiplier(base):
  if base == 'BTC':
    return 8000
  if base == 'BCH':
    return 1200
  elif base == 'ETH':
    return 600
  elif base == 'LTC':
    return 200
  elif base == 'EOS':
    return 16
  elif base == 'XRP':
    return 0.5
  elif base == 'FT':
    return 0.2
  elif base == 'USDT':
    return 1
  elif base == 'USD':
    return 1
  elif base == 'KRW':
    return 1 / 1000
  elif base == 'JPY':
    return 1 / 100
  else:
    raise ValueError(base)


def get_large_qty(product):
  if is_inverse(product):
    multiplier = get_futures_multiplier('Futures', product.exchange, product.symbol)
    return 5000 / multiplier

  price = get_multiplier(product.base) / get_multiplier(product.quote)
  return 8000 / price


def plot_table(ax, df):
  cellDict = ax.table(cellText=df.values, colLabels=df.columns, cellLoc='center',
                      loc='center').get_celld()
  for i in range(len(df.index) + 1):
    cellDict[(i, 0)].set_width(0.06)
    for j in range(1, len(df.columns)):
      cellDict[(i, j)].set_width(0.05)


def linecross(s1, t1, s2, t2):
  if s1 <= s2 <= t1 or s1 <= t2 <= t1:
    return True
  elif s2 <= s1 <= t2 or s2 <= t1 <= t2:
    return True
  else:
    return False


EventRecord = collections.namedtuple(
    'EventRecord', ['timestamp', 'type', 'bid_price', 'bid_qty', 'ask_price', 'ask_qty'])


class KuonaInteractive(object):
  def __init__(self, symbol, large_order_qty):
    self.symbol = symbol
    self.window_size = 128
    self.large_order_qty = large_order_qty
    self.cursor = 0
    self.rows = []
    self.obs = []
    self.evt_tss = []
    self.evts = []
    self.tss = []
    self.ask0s = []
    self.bid0s = []
    self.trd_tpq = []
    self.self_obs = []
    self.self_tss = []
    self.sub_tpq = []
    self.acc_tpq = []
    self.fill_tpq = []
    self.can_tpq = []
    self.fig = None
    self.die = False
    self.skip = False
    self.unit_step = 32
    self.iterate_until_cross = False
    self.plot_fill = True
    self.plot_sub = True
    self.plot_can = True
    self.plot_trd = True
    self.crsbook = []

  def md_cross_on_ob(self):
    if len(self.self_tss) == 0 or len(self.tss) == 0:
      return True
    elif linecross(self.tss[0], self.tss[-1], self.self_tss[0], self.self_tss[-1]):
      return True
    else:
      return False

  def get_events(self, n):
    import bisect
    begin_tss = self.tss[self.cursor]
    begin_idx = bisect.bisect_left(self.evt_tss, begin_tss)
    subdf = pandas.DataFrame(self.evts[begin_idx:(begin_idx + n)])
    subdf['timestamp'] = pandas.DatetimeIndex(subdf['timestamp']).strftime('%H:%M:%S.%f')
    return subdf

  def plot_canvas(self):
    if self.die:
      return
    tss = self.tss[self.cursor:(self.cursor + self.window_size)]
    obs = self.obs[self.cursor:(self.cursor + self.window_size)]
    ask0s = self.ask0s[self.cursor:(self.cursor + self.window_size)]
    bid0s = self.bid0s[self.cursor:(self.cursor + self.window_size)]

    updates, self_updates = putil.to_updates_and_plot_bbo(
        self.ax, tss, obs, ask0s, bid0s, self.trd_tpq if self.plot_trd else [], self.self_tss,
        self.self_obs, self.fill_tpq if self.plot_fill else [],
        self.sub_tpq if self.plot_sub else [], self.acc_tpq if self.plot_sub else [],
        self.can_tpq if self.plot_can else [])

    self.ax.set_title(f"{FLAGS.subsymbols} {FLAGS.title}")

    putil.plot_price_levels(self.ax,
                            updates,
                            putil.get_colorfunc_blue(max_threshold=self.large_order_qty,
                                                     min_threshold=self.large_order_qty * 1e-3),
                            lw=1)

    if len(self.crsbook) > 0:
      putil.plot_cross(self.ax, tss, self.crsbook)

    if FLAGS.exchange != 'Bitflyer':
      putil.plot_price_levels(self.ax, self_updates, lambda sdet: (1, 0, 1), linestyles='dashed')

    timestr = pandas.DatetimeIndex([tss[0]]).strftime('%Y%m%d %H:%M:%S')[0]
    plt.title("%s %s" % (self.symbol, timestr))
    plt.grid(False)
    plt.axis('on')

    evtdf = self.get_events(35)
    evtdf = evtdf.round(decimals=4)
    table = self.ax2.table(cellText=evtdf.values,
                           colLabels=evtdf.columns,
                           loc='center',
                           colWidths=[0.25, 0.08, 0.15, 0.15, 0.15, 0.15])
    table.set_fontsize(9)
    table.auto_set_font_size(False)
    # plot_table(self.ax2, evtdf)

  def press(self, event):
    if event.key == 'right':
      self.cursor += self.unit_step
    elif event.key == 'left':
      if self.cursor >= self.unit_step:
        self.cursor -= self.unit_step
    elif event.key == '1':
      self.plot_fill = not self.plot_fill
    elif event.key == '2':
      self.plot_sub = not self.plot_sub
    elif event.key == '3':
      self.plot_can = not self.plot_can
    elif event.key == '4':
      self.plot_trd = not self.plot_trd
    elif event.key == 'q':
      self.die = True
    elif event.key == 'n':
      self.skip = True
    elif event.key == '-':
      self.window_size *= 2
      self.unit_step *= 2
    elif event.key == '+':
      if self.unit_step > 1:
        self.window_size = int(self.window_size / 2)
        self.unit_step = int(self.unit_step / 2)
    elif event.key == '*':
      if len(self.self_tss) > 0:
        self.iterate_until_cross = True
    else:
      return

  def try_plot(self):
    while True:
      if self.skip:
        plt.close()
        return False
      if self.die:
        plt.close()
        sys.exit(1)
        return False
      if 0 <= self.cursor and self.cursor + self.window_size < len(self.tss):
        tss = self.tss[self.cursor:(self.cursor + self.window_size)]
        if self.iterate_until_cross:
          if not linecross(tss[0], tss[-1], self.self_tss[0], self.self_tss[-1]):
            self.cursor += self.unit_step
            continue
          else:
            self.iterate_until_cross = False

        if self.fig is None:
          self.fig = plt.figure(figsize=(17, 8))
          self.ax = plt.subplot2grid((1, 5), (0, 0), colspan=3, fig=self.fig)
          self.ax2 = plt.subplot2grid((1, 5), (0, 3), colspan=2, fig=self.fig)
          self.plot_canvas()
          self.fig.canvas.mpl_connect('key_press_event', self.press)
          plt.waitforbuttonpress(timeout=0)
        else:
          self.fig.clf()
          self.ax = plt.subplot2grid((1, 5), (0, 0), colspan=3, fig=self.fig)
          self.ax2 = plt.subplot2grid((1, 5), (0, 3), colspan=2, fig=self.fig)
          self.plot_canvas()
          plt.draw()
          plt.waitforbuttonpress(timeout=0)
      else:
        break


class Kuona(object):
  def __init__(self,
               order_product,
               feed_product,
               large_order_qty,
               num_levels=5,
               feed_product2=None,
               start_ts=None,
               end_ts=None):
    self.order_product = order_product
    self.feed_product = feed_product
    self.feed_product2 = feed_product2
    self.num_levels = num_levels
    self.order_manager = KuonaOrderManager(FLAGS.exchange, self.order_product)
    self.kuona_cursor = KuonaInteractive(symbol=feed_product.symbol,
                                         large_order_qty=large_order_qty)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self.feed_product, self.feed_product2]:
      if product is None:
        continue
      print(book_builder_name, book_builder)
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def from_fill_csv(self, fill_csv_file):
    df = pandas.read_csv(fill_csv_file)
    for _, row in df.iterrows():
      if row['type'] == 'FILL_DUMP':
        self.kuona_cursor.fill_tpq.append((row['time'], row['price'], row['qty'] * row['side']))
      elif row['type'] == 'SUBMIT_DUMP':
        self.kuona_cursor.sub_tpq.append((row['time'], row['price'], row['qty'] * row['side']))
      elif row['type'] == 'CANCEL_DUMP':
        self.kuona_cursor.can_tpq.append((row['time'], row['price'], row['qty'] * row['side']))
      else:
        import pdb
        pdb.set_trace()
        raise ValueError(row)

  def from_bins(self, tssbin, prcbin, sgnbin, dotbin):
    for tss, prc, sgn, dot in zip(tssbin, prcbin, sgnbin, dotbin):
      if dot == '4':
        self.kuona_cursor.sub_tpq.append((tss, prc, sgn))
      elif dot == '|':
        self.kuona_cursor.can_tpq.append((tss, prc, sgn))
      elif dot == 'o':
        self.kuona_cursor.fill_tpq.append((tss, prc, sgn))

  def from_excel_file(self, excel_file):
    exceldata1 = pandas.read_excel(excel_file, sheet_name='LimitOrder-1', header=1)
    exceldata2 = pandas.read_excel(excel_file, sheet_name='LimitOrder-2', header=1)

    df = exceldata1.append(exceldata2, ignore_index=True)
    print(df.groupby(['Contracts', 'Price type']).sum())

    cols = [
        column for column in df.columns
        if column.find('Cont') >= 0 or column.find('PL') >= 0 or column.find('Fee') >= 0
    ]
    df2 = df[cols]

    print(df2.groupby('Contracts').sum())

    df = df.loc[df['Contracts'] == self.order_product.native_symbol].reset_index(drop=True)
    df = df[::-1]

    df['price'] = df['Price($)']
    df['fill_qty'] = df['Filled(Cont)']
    df['qty'] = df['Amount(Cont)']
    df['side'] = ((df['Type'].str.find("Long") >= 0) - 0.5) * 2 * (
        (df['Type'].str.find("Open") >= 0) - 0.5) * 2
    df['time'] = pandas.DatetimeIndex(df['Date']).astype(int) - 8 * 3600 * 1e9
    prev_side_time = {}
    for _, row in df.iterrows():
      dtrow = row['time']
      if row['side'] in prev_side_time:
        self.kuona_cursor.sub_tpq.append(
            (prev_side_time[row['side']], row['price'], row['qty'] * row['side']))
      prev_side_time[row['side']] = row['time']
      self.kuona_cursor.can_tpq.append((dtrow, row['price'], row['qty'] * row['side']))
      if row['fill_qty'] > 0:
        self.kuona_cursor.fill_tpq.append((dtrow, row['price'], row['fill_qty'] * row['side']))

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      if not event.HasField('internal_order_id'):
        event.internal_order_id = str(event.proc_order_id)
      if FLAGS.exchange in ['Bitflyer', 'Bitmex']:
        pass
      else:
        if event.internal_order_id == '':
          return
      symbol = event.symbol
      if (symbol != self.order_product and symbol != self.order_product.native_symbol
          and symbol != self.order_product.symbol):
        return
      timestamp = event.event_time
      price = event.order_price
      qty = event.order_qty
      side = event.order_side
      event_type = event.type
      fill_price = event.fill_price
      fill_qty = event.fill_qty
      sign = get_order_sign(event.order_side)

      if event.type == OrderEvent.ORDER_SUBMITTED:
        self.kuona_cursor.sub_tpq.append((timestamp, price, qty * sign))

      if event.type == OrderEvent.ORDER_ACCEPTED:
        self.kuona_cursor.acc_tpq.append((timestamp, price, qty * sign))

      if event.type == OrderEvent.ORDER_FILLED:
        self.kuona_cursor.fill_tpq.append((timestamp, fill_price, fill_qty * sign))

      if (FLAGS.exchange == 'Bitflyer'
          and event.type == OrderEvent.CANCEL_SUBMITTED) or (event.type
                                                             == OrderEvent.CANCEL_CONFIRMED):
        if (event.external_order_id == ''
            and event.internal_order_id in self.order_manager.internal_to_external_map):
          event.external_order_id = self.order_manager.internal_to_external_map[
              event.internal_order_id]
        if event.external_order_id in self.order_manager.order_tpqs:
          tqps = self.order_manager.order_tpqs[event.external_order_id]
          price = tqps.price
          sign = tqps.sign
          qty = tqps.qty
          diff_qty = -qty
        self.kuona_cursor.can_tpq.append((timestamp, price, qty * sign))

      if self.order_manager.on_event(event):
        self.kuona_cursor.self_obs.append(dict(self.order_manager.self_order_book).copy())
        self.kuona_cursor.self_tss.append(timestamp)

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < trade.timestamp:
      return

    if self.kuona_cursor.skip:
      return
    if product != self.feed_product:
      return
    if trade.side == TradeSide.TRADE_BUY_SIDE:
      trade_qty = trade.qty
    elif trade.side == TradeSide.TRADE_SELL_SIDE:
      trade_qty = -trade.qty
    else:
      raise ValueError(trade.side)
    self.kuona_cursor.trd_tpq.append((trade.timestamp, trade.price, trade_qty))
    self.kuona_cursor.evt_tss.append(trade.timestamp)
    if trade_qty > 0:
      self.kuona_cursor.evts.append(EventRecord(trade.timestamp, 'T', 0, 0, trade.price,
                                                -trade.qty))
    else:
      self.kuona_cursor.evts.append(EventRecord(trade.timestamp, 'T', trade.price, -trade.qty, 0,
                                                0))

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    if product != self.feed_product:
      self.kuona_cursor.crsbook.append({
          'timestamp': book.timestamp,
          'bid0p': book.bid0().price,
          'ask0p': book.ask0().price,
      })
      return

    if not hasattr(self, 'last_book_time'):
      self.last_book_time = 0
    # if self.last_book_time + 2e8 > book.timestamp:
    #  return
    self.last_book_time = book.timestamp

    if self.kuona_cursor.skip:
      return

    if not book.has_bid():
      logging.info("no bid levels")
      return
    if not book.has_ask():
      logging.info("no ask levels")
      return
    if book.ask0().price <= book.bid0().price:
      return

    fields = collections.OrderedDict()
    n = self.num_levels

    bnames = ["bid%d" % (n - i - 1) for i in range(n)] + \
        ["ask%d" % i for i in range(n)]
    bnamemap = {}
    for i, bname in enumerate(bnames):
      bnamemap[bname] = "%s_%s" % (i, bname)

    # timestr = pandas.DatetimeIndex([book.timestamp]).strftime('%Y%m%d %H:%M:%S')[0]
    # fields['time'] = timestr
    fields['time'] = book.timestamp

    orderbook = {}
    if FLAGS.use_feed_cache:
      for side, method in [("bid", book.iter_bid), ("ask", book.iter_ask)]:
        for i, price_qty in enumerate(method()):
          if i >= n:
            break
          iname = "%s%d" % (side, i)
          name = bnamemap[iname]
          fields[iname + 'p'] = price_qty.price
          fields[iname + 'q'] = price_qty.qty
          orderbook[price_qty.price] = price_qty.qty
    else:
      for side, method in [("bid", book.iteritems_bid), ("ask", book.iteritems_ask)]:
        for i, (price, price_qty) in enumerate(method()):
          if i >= n:
            break
          iname = "%s%d" % (side, i)
          name = bnamemap[iname]
          fields[iname + 'p'] = price
          fields[iname + 'q'] = price_qty.qty
          orderbook[price] = price_qty.qty

    self.kuona_cursor.rows.append(fields)
    self.kuona_cursor.obs.append(orderbook)
    self.kuona_cursor.tss.append(book.timestamp)
    self.kuona_cursor.ask0s.append(fields['ask0p'])
    self.kuona_cursor.bid0s.append(fields['bid0p'])
    add_event = True
    if len(self.kuona_cursor.evts) > 0:
      last_event = self.kuona_cursor.evts[-1]
      if (last_event.type,
          last_event.bid_price,
          last_event.bid_qty,
          last_event.ask_price,
          last_event.ask_qty) == ('B',
                                  book.bid0().price,
                                  book.bid0().qty,
                                  book.ask0().price,
                                  book.ask0().qty):
        add_event = False
    if add_event:
      self.kuona_cursor.evt_tss.append(book.timestamp)
      self.kuona_cursor.evts.append(
          EventRecord(book.timestamp,
                      'B',
                      book.bid0().price,
                      book.bid0().qty,
                      book.ask0().price,
                      book.ask0().qty))
    if self.kuona_cursor.md_cross_on_ob():
      self.kuona_cursor.try_plot()


def launch(baseflags,
           fsflags,
           osflags,
           *,
           subsymbols,
           need_order=False,
           fill_csv_file=None,
           excel_file=None,
           pickle_file=None):
  products = abase.subsymbols_to_products(subsymbols, baseflags)

  baseflags.exchange, baseflags.market_type = subsymbols[0].split(":")[:2]

  start_time, end_time = abase.get_time_range(baseflags)

  strategy = Kuona(products[0],
                   products[0],
                   get_large_qty(products[0]),
                   flags.FLAGS.num_level,
                   products[1] if len(products) > 1 else None,
                   start_ts=start_time,
                   end_ts=end_time)

  print("order loading")
  if fill_csv_file:
    strategy.from_fill_csv(fill_csv_file)
  elif excel_file:
    strategy.from_excel_file(excel_file)
  elif pickle_file:
    import coin.util.slurm.base as slmbase
    sim_rows = slmbase.pickle_load_from_file(pickle_file)
    strategy.from_bins(**(sim_rows[0][-1]))
  elif osflags.HasField('orderlog_machine'):
    print(baseflags, osflags)
    abase.run_from_order_archive_base(baseflags, osflags, on_log=strategy.on_log)
  print("order loading done")

  if len(strategy.kuona_cursor.self_tss) == 0 and need_order:
    return

  print("feed loading")
  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)
  return 0


# ./pyrunner coin/strategy/mm/tool/kuona.py --trading_date=20190924 --subsymbols=Bitmex:Futures:BTC-USD.PERPETUAL --need_order=False --feed_machine=feed-01.ap-northeast-1.aws --use_feed_cache=True --time_range=0-2


def main(_):
  matplotlib.rcParams.update({'font.size': 7})
  # matplotlib.style.use('dark_background')
  # matplotlib.style.use('ggplot')
  matplotlib.style.use('seaborn')
  matplotlib.rcParams.update({'text.color': "blue", 'axes.labelcolor': "blue"})

  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()

  subsymbols = None if flags.FLAGS.subsymbols is None else flags.FLAGS.subsymbols.split(",")

  launch(baseflags,
         fsflags,
         osflags,
         subsymbols=subsymbols,
         need_order=flags.FLAGS.need_order,
         fill_csv_file=flags.FLAGS.fill_csv_file,
         excel_file=flags.FLAGS.excel_file,
         pickle_file=flags.FLAGS.pickle_file)


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  abase.define_order_archive_flags()

  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'

  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_bool('need_order', None, '')
  flags.DEFINE_string('fill_csv_file', None, 'fill_csv_file mode if given')
  flags.DEFINE_string('excel_file', None, 'excel_file mode if given')
  flags.DEFINE_string('pickle_file', None, '')
  flags.DEFINE_string('title', '', '')
  flags.DEFINE_integer('num_level', 5, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
